package org.example.iotdb.flink.sink;

import org.apache.flink.configuration.Configuration;
import org.apache.flink.streaming.api.functions.sink.RichSinkFunction;

import org.apache.iotdb.flink.Event;
import org.apache.iotdb.rpc.IoTDBConnectionException;
import org.apache.iotdb.rpc.StatementExecutionException;
import org.apache.iotdb.tsfile.file.metadata.enums.TSDataType;
import org.example.iotdb.session.pool.SessionPool;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.util.ArrayList;
import java.util.Arrays;
import java.util.List;

/**
 * @Auther: kevin
 * @Description:
 * @Company: 上海博般数据技术有限公司
 * @Version: 1.0.0
 * @Date: Created in 10:12 2021/8/20
 * @ProjectName: iotdb-learn
 */
public class IotDBSink extends RichSinkFunction<List<Event>> {

    private Logger logger = LoggerFactory.getLogger(this.getClass());

    private static transient SessionPool pool;

    @Override
    public void open(Configuration parameters) throws Exception {
        super.open(parameters);
        List<String> nodeUrls = Arrays.asList(
                "10.172.246.203:6667",
                "10.172.246.204:6667",
                "10.172.246.205:6667"
        );
        String user = "root";
        String password = "root";
        pool = new SessionPool(nodeUrls, user, password);
    }

    @Override
    public void invoke(List<Event> events, Context context) throws Exception {
        try {
            long start = System.currentTimeMillis();
            List<String> deviceIds = new ArrayList<>();
            List<Long> timestamps = new ArrayList<>();
            List<List<String>> measurementsList = new ArrayList<>();
            List<List<TSDataType>> typesList = new ArrayList<>();
            List<List<Object>> valuesList = new ArrayList<>();
            int cnt = 0;
            for (Event event : events) {
                deviceIds.add(event.getDevice());
                timestamps.add(event.getTimestamp());
                measurementsList.add(event.getMeasurements());
                typesList.add(event.getTypes());
                valuesList.add(event.getValues());
                cnt++;
                if (cnt % 200 == 0) {
                    pool.insertRecords(deviceIds, timestamps, measurementsList, typesList, valuesList);
                    deviceIds.clear();
                    timestamps.clear();
                    measurementsList.clear();
                    typesList.clear();
                    valuesList.clear();
                }
            }
            if (deviceIds.size() > 0) {
                pool.insertRecords(deviceIds, timestamps, measurementsList, typesList, valuesList);
            }
            long end = System.currentTimeMillis();
            logger.debug(String.format("thread[%s] batch insert %s record to iodTB. duration[%s]ms",
                    Thread.currentThread().getId(),
                    events.size(),
                    (end - start)));
        } catch (Exception e) {
            e.printStackTrace();
            logger.error("save to iotDB failed ! Cause:" + e.getMessage(), e);
        }
    }

    @Override
    public void close() throws Exception {
        pool.close();
    }
}
